import { z } from "zod";

const promptType = {
  SUMMARY: "summary",
  TREE_INSERT: "insert",
  TREE_SELECT: "tree_select",
  TREE_SELECT_MULTIPLE: "tree_select_multiple",
  QUESTION_ANSWER: "text_qa",
  REFINE: "refine",
  KEYWORD_EXTRACT: "keyword_extract",
  QUERY_KEYWORD_EXTRACT: "query_keyword_extract",
  SCHEMA_EXTRACT: "schema_extract",
  TEXT_TO_SQL: "text_to_sql",
  TEXT_TO_GRAPH_QUERY: "text_to_graph_query",
  TABLE_CONTEXT: "table_context",
  KNOWLEDGE_TRIPLET_EXTRACT: "knowledge_triplet_extract",
  SIMPLE_INPUT: "simple_input",
  PANDAS: "pandas",
  JSON_PATH: "json_path",
  SINGLE_SELECT: "single_select",
  MULTI_SELECT: "multi_select",
  VECTOR_STORE_QUERY: "vector_store_query",
  SUB_QUESTION: "sub_question",
  SQL_RESPONSE_SYNTHESIS: "sql_response_synthesis",
  SQL_RESPONSE_SYNTHESIS_V2: "sql_response_synthesis_v2",
  CONVERSATION: "conversation",
  DECOMPOSE: "decompose",
  CHOICE_SELECT: "choice_select",
  CUSTOM: "custom",
  RANKGPT_RERANK: "rankgpt_rerank",
} as const;

const promptTypeSchema = z.enum([
  promptType.SUMMARY,
  promptType.TREE_INSERT,
  promptType.TREE_SELECT,
  promptType.TREE_SELECT_MULTIPLE,
  promptType.QUESTION_ANSWER,
  promptType.REFINE,
  promptType.KEYWORD_EXTRACT,
  promptType.QUERY_KEYWORD_EXTRACT,
  promptType.SCHEMA_EXTRACT,
  promptType.TEXT_TO_SQL,
  promptType.TEXT_TO_GRAPH_QUERY,
  promptType.TABLE_CONTEXT,
  promptType.KNOWLEDGE_TRIPLET_EXTRACT,
  promptType.SIMPLE_INPUT,
  promptType.PANDAS,
  promptType.JSON_PATH,
  promptType.SINGLE_SELECT,
  promptType.MULTI_SELECT,
  promptType.VECTOR_STORE_QUERY,
  promptType.SUB_QUESTION,
  promptType.SQL_RESPONSE_SYNTHESIS,
  promptType.SQL_RESPONSE_SYNTHESIS_V2,
  promptType.CONVERSATION,
  promptType.DECOMPOSE,
  promptType.CHOICE_SELECT,
  promptType.CUSTOM,
  promptType.RANKGPT_RERANK,
]);

export const PromptType = promptTypeSchema.enum;
export type PromptType = z.infer<typeof promptTypeSchema>;
